mod convert;
use convert::{as_4u8_from_u32, as_u32_from4u8, as_u64_from8u8};

use std::fs::{File, OpenOptions};
use std::io::{Error, Read, Write};
use std::net::{TcpListener, TcpStream};
use std::path::Path;
use std::thread;

use clap::Parser;
use indicatif::{ProgressBar, ProgressStyle};

#[derive(Parser)]
#[command(name = "P2PTransFile")]
#[command(version = "1.0")]
#[command(
    about = "Transfer file by P2P",
    long_about = "Tranfer file by P2P, support client and server mode"
)]
struct Args {
    #[arg(short, long, default_value_t = false)]
    client: bool,
    #[arg(short, long, default_value_t = false)]
    sender: bool,
    #[arg(short, long, default_value_t = String::from("0.0.0.0"))]
    ip: String,
    #[arg(short, long, default_value_t = 9999, value_parser = clap::value_parser!(u16).range(1..))]
    port: u16,
    #[arg(short, long, required = true)]
    file: Vec<String>,
}

fn recv_data(stream: &mut TcpStream, buf: &mut [u8]) -> Result<usize, Error> {
    // recv 4B data
    stream.read_exact(&mut buf[0..4])?;
    let dlen = as_u32_from4u8(&buf[0..4]) as usize;
    // recv dlen data
    stream.read_exact(&mut buf[0..dlen])?;

    Ok(dlen)
}

fn send_data(stream: &mut TcpStream, buf: &[u8]) -> Result<(), Error> {
    // get data len
    let dlen = buf.len();
    let bytes_send = as_4u8_from_u32(dlen as u32);
    // send package header
    stream.write(&bytes_send)?;
    // send package
    stream.write(&buf)?;

    Ok(())
}

fn send_file(stream: &mut TcpStream, files: &Vec<String>) -> Result<(), Error> {
    for file in files {
        stream.write(&[1])?; // continue

        let path = Path::new(&file);
        if !path.exists() || !path.is_file() {
            panic!("path {} is not exist or dir", file);
        }

        let name = path.file_name().unwrap().to_str().unwrap();
        let flen = path.metadata().unwrap().len();

        println!(
            "File name: {} length: {:.03}MB",
            name,
            (flen as f32) / 1024.0 / 1024.0
        );

        send_data(stream, name.as_bytes())?;
        send_data(stream, &flen.to_be_bytes())?;

        const BUF_SIZE: usize = 65536; // 64KB

        let flen = path.metadata()?.len();

        let mut fpts = 0;

        let mut buf: [u8; BUF_SIZE] = [0; BUF_SIZE];
        let mut file = File::open(path).unwrap();

        let bar = ProgressBar::new(flen).with_style(
            ProgressStyle::with_template(
                "{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})",
            )
            .unwrap(),
        );

        loop {
            let dlen = file.read(&mut buf)?;

            send_data(stream, &buf[0..dlen])?;

            fpts += dlen as u64;

            bar.set_position(fpts);
            if fpts >= flen {
                break;
            }
        }
    }
    stream.write(&[0])?; // exit
    Ok(())
}

fn recv_file(stream: &mut TcpStream, save: &String) -> Result<(), Error> {
    let path = Path::new(&save);
    if !path.exists() || !path.is_dir() {
        panic!("path {} is not exist or dir", save);
    }

    let mut stage = [0; 1];
    let mut size_buf: [u8; 8] = [0; 8];
    let mut name_buf: [u8; 1024] = [0; 1024];
    let mut file_buf: [u8; 65536] = [0; 65536];

    loop {
        stream.read_exact(&mut stage)?;
        if stage[0] == 0 {
            break;
        }

        let dlen = recv_data(stream, &mut name_buf)?;
        let name = save.clone() + "/" + &String::from_utf8_lossy(&name_buf[0..dlen]).to_string();

        // let path = Path::new(&name);
        let dlen = recv_data(stream, &mut size_buf)?;
        assert_eq!(dlen, 8);

        let flen = as_u64_from8u8(&size_buf);

        println!(
            "File name: {} length: {:.03}MB",
            name,
            (flen as f32) / 1024.0 / 1024.0
        );

        let path = Path::new(&name);

        let mut file = OpenOptions::new().write(true).create(true).open(path)?;

        let mut dlen;
        let mut fpts = 0;

        let bar = ProgressBar::new(flen).with_style(
            ProgressStyle::with_template(
                "{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})",
            )
            .unwrap(),
        );

        loop {
            dlen = recv_data(stream, &mut file_buf)?;

            file.write(&file_buf[0..dlen])?;

            fpts += dlen as u64;

            bar.set_position(fpts);

            if fpts >= flen {
                break;
            }
        }
    }
    Ok(())
}

fn main() -> std::io::Result<()> {
    let _args = Args::parse();

    let as_server = !_args.client;
    let as_sender = _args.sender;

    let file_list = _args.file;

    assert!(file_list.len() > 0);

    let ipaddrs = _args.ip + ":" + &_args.port.to_string();

    if as_server {
        let listener = TcpListener::bind(&ipaddrs.trim())?;

        let mut thread_vec: Vec<thread::JoinHandle<()>> = Vec::new();

        for stream in listener.incoming() {
            let mut stream = stream.expect("failed!");

            let mut ls_file = Vec::new();
            for i in &file_list {
                ls_file.push(i.clone());
            }

            let handle = thread::spawn(move || {
                if as_sender {
                    send_file(&mut stream, &ls_file).unwrap_or_else(|error| eprint!("{:?}", error));
                } else {
                    recv_file(&mut stream, &ls_file[0])
                        .unwrap_or_else(|error| eprint!("{:?}", error));
                }
            });

            thread_vec.push(handle);
        }

        for handle in thread_vec {
            handle.join().unwrap();
        }
    } else {
        let mut stream = TcpStream::connect(&ipaddrs.trim())?;
        if as_sender {
            send_file(&mut stream, &file_list)?;
        } else {
            recv_file(&mut stream, &file_list[0])?;
        }
    }
    Ok(())
}
